#pragma once

#define VIRTUAL_ADDRESS_BITS 48
#define VIRTUAL_ADDRESS_MASK ((((ULONG_PTR)1) << VIRTUAL_ADDRESS_BITS) - 1)

#define PTE_BASE          0xFFFFF68000000000UI64
#define PFN_DATA_BASE	  0xFFFFFA8000000000UI64

#define PTI_SHIFT 12
#define PDI_SHIFT 21
#define PPI_SHIFT 30
#define PXI_SHIFT 39

#define PTE_SHIFT 3

#define _HARDWARE_PTE_WORKING_SET_BITS  11

typedef struct _MMPTE {
	ULONGLONG Valid : 1;
	ULONGLONG Writable : 1;        // changed for MP version
	ULONGLONG Owner : 1;
	ULONGLONG WriteThrough : 1;
	ULONGLONG CacheDisable : 1;
	ULONGLONG Accessed : 1;
	ULONGLONG Dirty : 1;
	ULONGLONG LargePage : 1;
	ULONGLONG Global : 1;
	ULONGLONG CopyOnWrite : 1; // software field
	ULONGLONG Prototype : 1;   // software field
	ULONGLONG Write : 1;       // software field - MP change
	ULONGLONG PageFrameNumber : 28;
	ULONG64 reserved1 : 24 - (_HARDWARE_PTE_WORKING_SET_BITS + 1);
	ULONGLONG SoftwareWsIndex : _HARDWARE_PTE_WORKING_SET_BITS;
	ULONG64 NoExecute : 1;
} MMPTE, *PMMPTE;


typedef struct _MMPFNENTRY                 // 11 elements, 0x2 bytes (sizeof) 
{
	struct                                 // 5 elements, 0x1 bytes (sizeof)  
	{
		/*0x000*/         UINT8        PageLocation : 3;     // 0 BitPosition                   
		/*0x000*/         UINT8        WriteInProgress : 1;  // 3 BitPosition                   
		/*0x000*/         UINT8        Modified : 1;         // 4 BitPosition                   
		/*0x000*/         UINT8        ReadInProgress : 1;   // 5 BitPosition                   
		/*0x000*/         UINT8        CacheAttribute : 2;   // 6 BitPosition                   
	};
	struct                                 // 6 elements, 0x1 bytes (sizeof)  
	{
		/*0x001*/         UINT8        Priority : 3;         // 0 BitPosition                   
		/*0x001*/         UINT8        Rom : 1;              // 3 BitPosition                   
		/*0x001*/         UINT8        InPageError : 1;      // 4 BitPosition                   
		/*0x001*/         UINT8        KernelStack : 1;      // 5 BitPosition                   
		/*0x001*/         UINT8        RemovalRequested : 1; // 6 BitPosition                   
		/*0x001*/         UINT8        ParityError : 1;      // 7 BitPosition                   
	};
}MMPFNENTRY, *PMMPFNENTRY;

typedef struct _MMPFN                                // 13 elements, 0x30 bytes (sizeof) 
{
	union                                            // 7 elements, 0x8 bytes (sizeof)   
	{
		/*0x000*/         UINT64       Flink;
		/*0x000*/         ULONG32      WsIndex;
		/*0x000*/         struct _KEVENT* Event;
		/*0x000*/         VOID*        Next;
		/*0x000*/         VOID*        VolatileNext;
		/*0x000*/         struct _KTHREAD* KernelStackOwner;
		/*0x000*/         struct _SINGLE_LIST_ENTRY NextStackPfn;      // 1 elements, 0x8 bytes (sizeof)   
	}u1;
	union                                            // 3 elements, 0x8 bytes (sizeof)   
	{
		/*0x008*/         UINT64       Blink;
		/*0x008*/         struct _MMPTE* ImageProtoPte;
		/*0x008*/         UINT64       ShareCount;
	}u2;
	union                                            // 4 elements, 0x8 bytes (sizeof)   
	{
		/*0x010*/         struct _MMPTE* PteAddress;
		/*0x010*/         VOID*        VolatilePteAddress;
		/*0x010*/         LONG32       Lock;
		/*0x010*/         UINT64       PteLong;
	};
	union                                            // 3 elements, 0x4 bytes (sizeof)   
	{
		struct                                       // 2 elements, 0x4 bytes (sizeof)   
		{
			/*0x018*/             UINT16       ReferenceCount;
			/*0x01A*/             struct _MMPFNENTRY e1;                   // 11 elements, 0x2 bytes (sizeof)  
		};
		struct                                       // 3 elements, 0x4 bytes (sizeof)   
		{
			union                                    // 2 elements, 0x2 bytes (sizeof)   
			{
				/*0x018*/                 UINT16       ReferenceCount;
				/*0x018*/                 INT16        VolatileReferenceCount;
			};
			/*0x01A*/             UINT16       ShortFlags;
		}e2;
	}u3;
	/*0x01C*/     UINT16       UsedPageTableEntries;
	/*0x01E*/     UINT8        VaType;
	/*0x01F*/     UINT8        ViewCount;
	union                                            // 2 elements, 0x8 bytes (sizeof)   
	{
		/*0x020*/         struct _MMPTE OriginalPte;                   // 1 elements, 0x8 bytes (sizeof)   
		/*0x020*/         LONG32       AweReferenceCount;
	};
	union                                            // 6 elements, 0x8 bytes (sizeof)   
	{
		struct                                       // 6 elements, 0x8 bytes (sizeof)   
		{
			/*0x028*/             UINT64       PteFrame : 52;              // 0 BitPosition                    
			/*0x028*/             UINT64       Unused : 3;                 // 52 BitPosition                   
			/*0x028*/             UINT64       PfnImageVerified : 1;       // 55 BitPosition                   
			/*0x028*/             UINT64       AweAllocation : 1;          // 56 BitPosition                   
			/*0x028*/             UINT64       PrototypePte : 1;           // 57 BitPosition                   
			/*0x028*/             UINT64       PageColor : 6;              // 58 BitPosition                   
		};
	}u4;
}MMPFN, *PMMPFN;

#define MiGetPteAddress(va) \
    ((PMMPTE)(((((ULONG_PTR)(va) & VIRTUAL_ADDRESS_MASK) >> PTI_SHIFT) << PTE_SHIFT) + PTE_BASE))

typedef struct _MAP_STRUCT {
	PVOID OrigPage;
	PVOID RealPage;
	PVOID MapPage;
	PMDL Mdl;
	PHYSICAL_ADDRESS MapPagePhys;
} MAP_STRUCT, *PMAP_STRUCT;

PVOID
MapSpecifiedPage(
	IN PMAP_STRUCT MapHandle,
	IN PHYSICAL_ADDRESS PhysicalAddress
	);

VOID
FiniMapPage(
	IN PMAP_STRUCT MapHandle
	);

NTSTATUS
InitMapPage(
	OUT PMAP_STRUCT MapHandle
	);
